import torch.nn as nn
import torch.nn.functional as F
from transformers import BertConfig, BertModel
import logging
import json

logger = logging.getLogger(__name__)


class Bert(nn.Module):
    """基于BERT预训练模型的文本分类器"""

    def __init__(self, config):
        """
        初始化模型参数和结构

        Args:
            config: 包含模型配置的对象，需包含以下属性：
                - config_file: BERT配置文件路径
                - model_name_or_path: 预训练模型名称或路径
                - num_labels: 分类标签数量
                - task: 微调任务名称
                - requires_grad: 是否允许BERT参数更新
                - hidden_dropout_prob: Dropout概率
                - multi_drop: 多轮dropout的权重
                - hidden_size: BERT隐藏层维度
        """
        super(Bert, self).__init__()

        # 从配置文件读取BERT模型参数
        with open(config.config_file, 'r', encoding='utf-8') as f:
            config_dict = json.load(f)

        # 更新配置参数，添加分类任务相关设置
        config_dict['num_labels'] = config.num_labels
        config_dict['finetuning_task'] = config.task

        # 创建BERT配置对象并加载预训练模型
        model_config = BertConfig.from_dict(config_dict)
        self.bert = BertModel.from_pretrained(config.model_name_or_path, config=model_config)

        # 设置BERT参数是否参与训练
        if config.requires_grad:
            for param in self.bert.parameters():
                param.requires_grad = True

        # 定义dropout层和分类器
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.multi_drop = config.multi_drop
        self.hidden_size = config.hidden_size
        self.classifier = nn.Sequential(nn.Linear(self.hidden_size, config.num_labels))

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, labels=None):
        """
        模型前向传播过程

        Args:
            input_ids: 输入文本的token ID序列
            attention_mask: 注意力掩码，区分真实token和padding
            token_type_ids: 句子类型ID，区分两个句子（用于问答等任务）
            labels: 分类标签（训练时使用）

        Returns:
            out: 最终分类logits
            loss: 计算的损失值（如果提供了标签）
        """
        # 通过BERT模型获取输出
        outputs = self.bert(input_ids=input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids)

        # 获取[CLS]标记的表示作为句子整体表示
        pooled_output = outputs[1]  # [batch_size, hidden_size]

        # 初始化输出和损失
        out = None
        loss = 0

        # 多轮dropout机制：多次应用dropout和分类器，提高模型稳定性
        for i in range(5):
            # 应用dropout
            dropped_output = self.dropout(pooled_output)
            # 通过分类器获取logits
            current_out = self.classifier(dropped_output)

            # 累加输出结果（最后一轮的结果作为最终输出）
            out = current_out

            # 如果提供了标签，则计算损失
            if labels is not None:
                if i == 0:
                    # 第一轮计算的损失进行平均
                    loss = F.cross_entropy(current_out, labels) / self.multi_drop
                else:
                    # 累加后续轮次的损失
                    loss += F.cross_entropy(current_out, labels) / self.multi_drop

        return out, loss
